import os
import torch
import numpy as np
from torch import nn
from torchvision import transforms
from cvx2 import WidthBlock, DOWidthBlock, SE, CA, ECA, CBAM
from torchvision.datasets import ImageFolder
from cvx2.wrapper import SplitImageClassifyModelWrapper
from model_wrapper import log_utils
from cvx2.utils import get_pretrained, get_img_mean_std
from torchvision.models import resnet18, ResNet18_Weights

imgsz = 28
data_dir = '/Users/summy/data/AI-face'

if __name__ == '__main__':
	mean, std = get_img_mean_std(data_dir, imgsz=imgsz, sample=1000)
	print('mean:', mean, 'std:', std)
	transform = transforms.Compose([
		transforms.Resize((imgsz, imgsz)),  # 将所有图片resize到28x28
		transforms.ToTensor(),
		transforms.Normalize(mean=mean, std=std)
	])
	train_transform = transforms.Compose([
		transforms.Resize((imgsz, imgsz)),  # 将所有图片resize到28x28
		transforms.RandomAffine(degrees=(-10, 10), translate=(0.1, 0.1), scale=(0.9, 1.1), shear=(-10, 10)),
		transforms.ToTensor(),
		transforms.Normalize(mean=mean, std=std)
		
	])
	dataset = ImageFolder(root=data_dir, transform=transform)
	# dataset = ImageFolder(root=data_dir, transform=ResNet18_Weights.DEFAULT.transforms())
	classes = dataset.classes
	
	model = nn.Sequential(
		WidthBlock(c1=3, c2=32),
		nn.MaxPool2d(kernel_size=2, stride=2),
		WidthBlock(c1=32, c2=32),
		nn.MaxPool2d(kernel_size=2, stride=2),
		WidthBlock(c1=32, c2=32),
		# DOWidthBlock 68.74%
		# WidthBlock 68.26%
		# SE(32), # 64.45%
		# CA(32), # 63.00%
		# ECA(32),  # 64.90%
		CBAM(32), # 61.97%
		nn.Flatten(),
		nn.Linear(in_features=32 * (imgsz >> 2) ** 2, out_features=1024),
		nn.Dropout(0.2),
		nn.SiLU(inplace=True),
		nn.Linear(in_features=1024, out_features=len(classes)),
	)

	# model = get_pretrained(resnet18, ResNet18_Weights.DEFAULT, len(classes))
	wrapper = SplitImageClassifyModelWrapper(model)
	wrapper.train_evaluate(data=dataset, val_size=0.2, random_state=42, amp=False,
						train_transform=train_transform, epochs=1, batch_size=64, lr=0.0001)
	# print(wrapper.evaluate(data=data_dir, transform=transform, batch_size=64))
	# print(wrapper.classification_report(data=data_dir, transform=transform, batch_size=64))
	log_utils.info('跑完了')
